from typing import Dict, Any, Optional, List
from src.core.output_manager_base import OutputManagerBase, SimulationParameters, MPI, FENICSX_AVAILABLE
import csv
import os

class CSVDataLogger(OutputManagerBase):
    """
    用于将特定数据点、标量值或时间序列数据记录到 CSV 文件。
    """
    def __init__(self, 
                 comm: Optional[MPI.Comm], 
                 params: SimulationParameters, 
                 fem_entities: Dict[str, Any], 
                 output_dir: str,
                 csv_filename: str = "datalog.csv",
                 log_frequency: int = 1,
                 headers: Optional[List[str]] = None, # CSV 文件的表头
                 data_extraction_callback = None): # 一个回调函数，用于从 fem_entities 或 params 提取数据行
        """
        初始化 CSVDataLogger。

        参数:
            comm (Optional[MPI.Comm]): MPI 通信器。
            params (SimulationParameters): 模拟参数。
            fem_entities (Dict[str, Any]): FEM 实体。
            output_dir (str): 输出目录。
            csv_filename (str): CSV 文件名。
            log_frequency (int): 日志记录频率。
            headers (Optional[List[str]]): CSV 文件的表头。如果为 None，则可能由回调函数决定或不写入表头。
            data_extraction_callback (Callable[[int, float, Dict, Any], Optional[List[Any]]]): 
                一个回调函数，签名为 (iter_count, current_time, fem_entities, params) -> Optional[List[Any]]，
                返回要在当前步骤记录的数据行 (列表)。如果返回 None，则不记录。
        """
        super().__init__(comm, params, fem_entities, output_dir)
        self.csv_filename = csv_filename
        self.log_frequency = log_frequency
        self.headers = headers
        self.data_extraction_callback = data_extraction_callback
        
        self.csv_filepath = os.path.join(self.output_dir, self.csv_filename)
        self.file_handle = None
        self.csv_writer = None

        if self.comm is None or self.comm.rank == 0: # type: ignore
            print(f"CSVDataLogger: 初始化。日志频率: {self.log_frequency}, 文件: '{self.csv_filepath}'")
            self._open_csv_file() # 在主进程上打开文件并写入表头

    def _open_csv_file(self):
        """在主进程上打开 CSV 文件并准备写入器。"""
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return
        
        os.makedirs(self.output_dir, exist_ok=True)
        try:
            # 使用 'a' 模式以追加，如果需要从头开始则用 'w'
            # 但通常日志是追加的，除非在 finalize 中有特殊处理
            # 为简单起见，如果文件已存在，这里会覆盖。更好的做法是检查或追加。
            self.file_handle = open(self.csv_filepath, 'w', newline='')
            self.csv_writer = csv.writer(self.file_handle)
            if self.headers:
                self.csv_writer.writerow(self.headers)
                print(f"  CSV 文件表头已写入: {self.headers}")
        except IOError as e:
            print(f"错误: CSVDataLogger 无法打开 CSV 文件 '{self.csv_filepath}' 进行写入: {e}")
            self.file_handle = None
            self.csv_writer = None

    def is_output_step(self, iteration_count: int, current_time: float) -> bool:
        """
        判断是否为日志记录步骤。
        """
        # 数据提取和写入通常只在主进程执行，以避免文件冲突
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return False
        return iteration_count % self.log_frequency == 0

    def execute_step(self, iteration_count: int, current_time: float) -> None:
        """
        执行 CSV 日志记录步骤。
        """
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return

        if self.csv_writer is None:
            print("CSVDataLogger: CSV writer 未初始化，跳过日志记录。")
            return

        if self.data_extraction_callback:
            try:
                # 调用回调函数获取当前步骤的数据
                # (iter_count, current_time, fem_entities, params)
                row_data = self.data_extraction_callback(iteration_count, current_time, self.fem_entities, self.params)
                if row_data is not None:
                    if not isinstance(row_data, list):
                        print(f"警告: CSVDataLogger 的 data_extraction_callback 未返回列表 (而是 {type(row_data)})。跳过写入。")
                        return
                    self.csv_writer.writerow(row_data)
                    # print(f"  CSVDataLogger: 在迭代 {iteration_count} 记录数据: {row_data}")
                # else:
                    # print(f"  CSVDataLogger: 回调函数在迭代 {iteration_count} 未返回数据，跳过写入。")
            except Exception as e:
                print(f"错误: CSVDataLogger 执行 data_extraction_callback 时出错: {e}")
        else:
            # 如果没有回调，可以记录默认信息，例如迭代次数和时间
            default_row = [iteration_count, f"{current_time:.4e}"]
            self.csv_writer.writerow(default_row)
            if not self.headers: # 如果没有预设表头，第一次记录时尝试添加默认表头
                self.headers = ["Iteration", "Time"]
                # 这有点晚，但可以补救 (需要重新打开文件以 'w' 模式写入表头，然后追加)
                # 为简单起见，这里不处理这种情况。建议始终提供表头或通过回调处理。
                print(f"  CSVDataLogger: 记录默认数据 (迭代, 时间)。建议提供表头和回调。")

    def finalize(self) -> None:
        """
        完成 CSV 日志记录操作，例如关闭文件。
        """
        if self.comm is not None and self.comm.rank != 0: # type: ignore
            return
            
        if self.file_handle:
            try:
                self.file_handle.close()
                print(f"CSVDataLogger: CSV 文件 '{self.csv_filepath}' 已关闭。")
            except IOError as e:
                print(f"错误: CSVDataLogger 关闭 CSV 文件时出错: {e}")
        self.file_handle = None
        self.csv_writer = None
        print("CSVDataLogger: 完成所有日志操作。") 